import pandas as pd
from MyTT import *
import warnings
import numpy  as np
warnings.filterwarnings('ignore', 'invalid value encountered in divide')
# 读取原始数据
df = pd.read_parquet('G:/quant_data/merged_a_stock_data.parquet')
CLOSE = df.close.values
OPEN = df.open.values
HIGH = df.high.values
LOW = df.low.values
VOLUME = df.volume.values

# 计算各指标
# 定义计算CCI的函数
def calculate_cci(high, low, close, period=14):
    """
    计算CCI指标
    :param high: 高价数组 (numpy array)
    :param low: 低价数组 (numpy array)
    :param close: 收盘价数组 (numpy array)
    :param period: 计算周期，默认为14
    :return: CCI指标数组 (numpy array)
    """
    if len(high) != len(low) or len(low) != len(close):
        raise ValueError("high, low, close数组长度必须一致")
    
    # 计算典型价格 TP
    tp = (high + low + close) / 3.0
    
    # 初始化CCI数组
    cci = np.zeros(len(tp))
    
    # 计算每个周期的CCI
    for i in range(period - 1, len(tp)):
        # 当前周期的典型价格
        tp_period = tp[i - period + 1:i + 1]
        
        # 计算MA
        ma = np.mean(tp_period)
        
        # 计算MD
        md = np.mean(np.abs(tp_period - ma))
        
        # 计算CCI
        if md == 0:
            cci[i] = 0  # 避免除以零
        else:
            cci[i] = (tp[i] - ma) / (0.015 * md)
    
    return cci
# 计算LWRciliHn = HHV(HIGH, N=14)      # N周期最高价序列[3,6](@ref)
Ln = LLV(LOW, N=14)       # N周期最低价序列[3,6](@ref)
wr10, _ = WR(CLOSE, HIGH, LOW, N=10) #WR 威廉
df['WR10'] = wr10
_, wr6 = WR(CLOSE, HIGH, LOW, N=6) 
df['WR6'] = wr6 # 威廉指标 [6,10]
df['K'], df['D'], df['J'] = KDJ(CLOSE, HIGH, LOW)    # 随机指标
df['DIF'], df['DEA'], df['MACD'] = MACD(CLOSE)   #MACD 
df['RSI1'] = RSI(CLOSE,6)   
df['RSI2'] = RSI(CLOSE,12) 
df['RSI3'] = RSI(CLOSE,24)                               # 相对强弱指标 [15](@ref)
df['DMA'] = DMA(CLOSE,0.1)                                 # 动态移动平均 [6](@ref)
df['EXPMA'] = EMA(CLOSE,12)                               # 指数移动平均 [20](@ref)
#df['AVEDEV'] = AVEDEV(CLOSE,20)                           # 平均偏差 [21](@ref)
df['TTRR'] = ATR(CLOSE, HIGH, LOW,14)                     # 真实波幅 [3](@ref)
df['PDI'], df['MDI'], df['ADX'], _ = DMI(HIGH, LOW, CLOSE)  # 趋向指标 [6,12](@ref) 
df['MTM'] = CLOSE - REF(CLOSE, N=12)  # 动量差值[7,9](@ref)
df['ABS'] = ABS(CLOSE)
# 计算各周期均线
MA3 = MA(CLOSE, 3)    # 3日均线[7,9](@ref)
MA6 = MA(CLOSE, 6)    # 6日均线
MA12 = MA(CLOSE, 12)  # 12日均线
MA24 = MA(CLOSE, 24)  # 24日均线

# 计算BBI
df['BBI'] = (MA3 + MA6 + MA12 + MA24) / 4  # 多空均线合成[7,9](@ref)
#CCI 用numpy 计算
df['CCI'] = calculate_cci(HIGH, LOW, CLOSE, period=14)
# 确保关键字段存在（新增code字段检查）
if 'code' not in df.columns:
    df['code'] = 'default_code'  # 如果原始数据没有code字段，设置默认值

# 字段排序（date和code在前）
cols_order = ['date', 'code'] + [col for col in df.columns if col not in ['date', 'code']]
df['date'] = pd.to_datetime(df['date']).dt.normalize().dt.tz_localize(None)
df = df[cols_order]
df_rounded = df.round(2)
# ================= 保存结果 =================
output_path = 'G:/quant_data/basic_indicators.parquet'
df_rounded.to_parquet(output_path, index=False)